import or1200_definitions
import sys
sys.path.append("..")
import pyrtl
from  pyrtl import GPUSim

width = or1200_definitions.OR1200_OPERAND_WIDTH

class Or1200Operandmuxes(object):
    def __init__(self):
        self.rst = pyrtl.Input(bitwidth=1, name='operandmuxes_rst')
        self.id_freeze = pyrtl.Input(bitwidth=1, name='operandmuxes_id_freeze')
        self.ex_freeze = pyrtl.Input(bitwidth=1, name='operandmuxes_ex_freeze')
        self.rf_dataa = pyrtl.Input(bitwidth=width, name='operandmuxes_rf_dataa')
        self.rf_datab = pyrtl.Input(bitwidth=width, name='operandmuxes_rf_datab')
        self.ex_forw = pyrtl.Input(bitwidth=width, name='operandmuxes_ex_forw')
        self.wb_forw = pyrtl.Input(bitwidth=width, name='operandmuxes_wb_forw')
        self.simm = pyrtl.Input(bitwidth=width, name='operandmuxes_simm')
        self.sel_a = pyrtl.Input(bitwidth=or1200_definitions.OR1200_SEL_WIDTH, name='operandmuxes_sel_a')
        self.sel_b = pyrtl.Input(bitwidth=or1200_definitions.OR1200_SEL_WIDTH, name='operandmuxes_sel_b')
        self.operand_a = pyrtl.Output(bitwidth=width,name='operand_a ')
        self.operand_b = pyrtl.Output(bitwidth=width,name='operand_b')

        
        # self.operand_a = pyrtl.Register(bitwidth=width)
        # self.operand_b = pyrtl.Register(bitwidth=width)
        self.muxed_a = pyrtl.Register(bitwidth=width)
        self.muxed_b = pyrtl.Register(bitwidth=width)

        operareg = OperAReg()
        operbreg = OperBReg()
        forlogica = ForLogicA()
        forlogicb = ForLogicB()
        
        operareg.rst <<= self.rst
        self.operand_a <<= operareg.operand_a
        operareg.ex_freeze <<= self.ex_freeze
        operareg.id_freeze <<= self.id_freeze
        operareg.muxed_a <<= self.muxed_a
        operbreg.rst <<= self.rst
        self.operand_b <<= operbreg.operand_b
        operbreg.ex_freeze <<= self.ex_freeze
        operbreg.id_freeze <<= self.id_freeze
        operbreg.muxed_b <<= self.muxed_b
        forlogica.sel_a <<= self.sel_a
        self.muxed_a.next <<= forlogica.muxed_a
        forlogica.ex_forw <<= self.ex_forw
        forlogica.wb_forw <<= self.wb_forw
        forlogica.rf_dataa <<= self.rf_dataa
        forlogicb.sel_b <<= self.sel_b
        self.muxed_b.next <<= forlogicb.muxed_b
        forlogicb.ex_forw <<= self.ex_forw
        forlogicb.wb_forw <<= self.wb_forw
        forlogicb.rf_datab <<= self.rf_datab
        forlogicb.simm <<= self.simm


class Or1200Operandmuxes_wire(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1, name='operandmuxes_rst')
        self.id_freeze = pyrtl.WireVector(bitwidth=1, name='operandmuxes_id_freeze')
        self.ex_freeze = pyrtl.WireVector(bitwidth=1, name='operandmuxes_ex_freeze')
        self.rf_dataa = pyrtl.WireVector(bitwidth=width, name='operandmuxes_rf_dataa')
        self.rf_datab = pyrtl.WireVector(bitwidth=width, name='operandmuxes_rf_datab')
        self.ex_forw = pyrtl.WireVector(bitwidth=width, name='operandmuxes_ex_forw')
        self.wb_forw = pyrtl.WireVector(bitwidth=width, name='operandmuxes_wb_forw')
        self.simm = pyrtl.WireVector(bitwidth=width, name='operandmuxes_simm')
        self.sel_a = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_SEL_WIDTH, name='operandmuxes_sel_a')
        self.sel_b = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_SEL_WIDTH, name='operandmuxes_sel_b')
        self.operand_a = pyrtl.WireVector(bitwidth=width, name='operand_a ')
        self.operand_b = pyrtl.WireVector(bitwidth=width, name='operand_b')


        self.muxed_a = pyrtl.Register(bitwidth=width)
        self.muxed_b = pyrtl.Register(bitwidth=width)

        operareg = OperAReg()
        operbreg = OperBReg()
        forlogica = ForLogicA()
        forlogicb = ForLogicB()

        operareg.rst <<= self.rst
        self.operand_a <<= operareg.operand_a
        operareg.ex_freeze <<= self.ex_freeze
        operareg.id_freeze <<= self.id_freeze
        operareg.muxed_a <<= self.muxed_a
        operbreg.rst <<= self.rst
        self.operand_b <<= operbreg.operand_b
        operbreg.ex_freeze <<= self.ex_freeze
        operbreg.id_freeze <<= self.id_freeze
        operbreg.muxed_b <<= self.muxed_b
        forlogica.sel_a <<= self.sel_a
        self.muxed_a.next <<= forlogica.muxed_a
        forlogica.ex_forw <<= self.ex_forw
        forlogica.wb_forw <<= self.wb_forw
        forlogica.rf_dataa <<= self.rf_dataa
        forlogicb.sel_b <<= self.sel_b
        self.muxed_b.next <<= forlogicb.muxed_b
        forlogicb.ex_forw <<= self.ex_forw
        forlogicb.wb_forw <<= self.wb_forw
        forlogicb.rf_datab <<= self.rf_datab
        forlogicb.simm <<= self.simm
        
#Operand A register
class OperAReg(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1)
        self.operand_a = pyrtl.WireVector(bitwidth=32)
        self.saved_a = pyrtl.Register(bitwidth=1)
        self.ex_freeze = pyrtl.WireVector(bitwidth=1)
        self.id_freeze = pyrtl.WireVector(bitwidth=1)
        self.muxed_a = pyrtl.WireVector(bitwidth=32)
        
        with pyrtl.conditional_assignment:
            with self.rst == or1200_definitions.ONE:
                self.operand_a |= pyrtl.Const(0, bitwidth=32)
                self.saved_a.next |= pyrtl.Const(0b0, bitwidth=1)
            with ~self.ex_freeze == or1200_definitions.ONE:
                with self.id_freeze == or1200_definitions.ONE:
                    with ~self.saved_a == or1200_definitions.ONE:
                        self.operand_a |= self.muxed_a
                        self.saved_a.next |= pyrtl.Const(0b1, bitwidth=1)
            with ~self.ex_freeze == or1200_definitions.ONE:
                with ~self.saved_a == or1200_definitions.ONE:
                    self.operand_a |= self.muxed_a
            with ~self.ex_freeze == or1200_definitions.ONE:
                with ~self.id_freeze == or1200_definitions.ONE:
                    self.saved_a.next |= pyrtl.Const(0b0, bitwidth=1)
                
#Operand B register
class OperBReg(object):
    def __init__(self):
        self.rst = pyrtl.WireVector(bitwidth=1)
        self.operand_b = pyrtl.WireVector(bitwidth=32)
        self.saved_b = pyrtl.Register(bitwidth=1)
        self.ex_freeze = pyrtl.WireVector(bitwidth=1)
        self.id_freeze = pyrtl.WireVector(bitwidth=1)
        self.muxed_b = pyrtl.WireVector(bitwidth=32)
        
        with pyrtl.conditional_assignment:
            with self.rst == or1200_definitions.ONE:
                self.operand_b |= pyrtl.Const(0, bitwidth=32)
                self.saved_b.next |= pyrtl.Const(0b0, bitwidth=1)
            with ~self.ex_freeze == or1200_definitions.ONE:
                with self.id_freeze == or1200_definitions.ONE:
                    with ~self.saved_b == or1200_definitions.ONE:
                        self.operand_b |= self.muxed_b
                        self.saved_b.next |= pyrtl.Const(0b1, bitwidth=1)
            with ~self.ex_freeze == or1200_definitions.ONE:
                with ~self.saved_b == or1200_definitions.ONE:
                    self.operand_b |= self.muxed_b
            with ~self.ex_freeze == or1200_definitions.ONE:
                with ~self.id_freeze == or1200_definitions.ONE:
                    self.saved_b.next |= pyrtl.Const(0b0, bitwidth=1)
                
#Forwarding logic for operand A register
class ForLogicA(object):
    def __init__(self):
        self.sel_a = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_SEL_WIDTH)
        self.muxed_a = pyrtl.WireVector(bitwidth=32)
        self.ex_forw = pyrtl.WireVector(bitwidth=32)
        self.wb_forw = pyrtl.WireVector(bitwidth=32)
        self.rf_dataa = pyrtl.WireVector(bitwidth=32)
        
        with pyrtl.conditional_assignment:
            with self.sel_a == or1200_definitions.OR1200_SEL_EX_FORW:
                self.muxed_a |= self.ex_forw
            with self.sel_a == or1200_definitions.OR1200_SEL_WB_FORW:
                self.muxed_a |= self.wb_forw
            with pyrtl.otherwise:
                self.muxed_a |= self.rf_dataa
                
#Forwarding logic for operand A register
class ForLogicB(object):
    def __init__(self):
        self.sel_b = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_SEL_WIDTH)
        self.muxed_b = pyrtl.WireVector(bitwidth=32)
        self.ex_forw = pyrtl.WireVector(bitwidth=32)
        self.wb_forw = pyrtl.WireVector(bitwidth=32)
        self.rf_datab = pyrtl.WireVector(bitwidth=32)
        self.simm = pyrtl.WireVector(bitwidth=width)
        
        with pyrtl.conditional_assignment:
            with self.sel_b == or1200_definitions.OR1200_SEL_IMM:
                self.muxed_b |= self.simm
            with self.sel_b == or1200_definitions.OR1200_SEL_EX_FORW:
                self.muxed_b |= self.ex_forw
            with self.sel_b == or1200_definitions.OR1200_SEL_WB_FORW:
                self.muxed_b |= self.wb_forw
            with pyrtl.otherwise:
                self.muxed_b |= self.rf_datab
                
if __name__ == '__main__':
    operandmuxes = Or1200Operandmuxes()
    sim = pyrtl.GPUSim_now.GPUSim('oper')

    com = pyrtl.CompiledSimulation()

    input_width = []
    from pyrtl.wire import Input

    IN = pyrtl.working_block().wirevector_subset(Input)
    for inn in IN:
        name = inn.name
        width = inn.bitwidth
        input_width.append([name, width])
    import random

    max_ = 0
    com_time = 0
    ess_time = 0
    num = 1
    for i in range(num):
        print(i)

        inputs_random = {}
        inputs_select = {}
        inputs_old = {}
        inputs_ = {}

        for inp in input_width:
            inputs_random[inp[0]] = []
            inputs_select[inp[0]] = []
            result = random.randint(0, 2 ** inp[1] - 1)
            inputs_old[inp[0]] = result

        for j in range(131072):
            for inp in input_width:
                result = random.randint(0, 2 ** inp[1] - 1)
                inputs_[inp[0]] = result
                inputs_random[inp[0]].append(result)
            if random.random() < 0.2:
                for name in inputs_old:
                    inputs_select[name].append(inputs_old[name])
            else:
                for name in inputs_:
                    inputs_select[name].append(inputs_[name])
                    inputs_old[name] = inputs_[name]
        time1 = com.step_multiple(inputs_select)

    sim.make_input(inputs_select)